from fastapi import FastAPI, UploadFile, File, APIRouter
from pydantic import BaseModel
from typing import List
import cv2
from paddleocr import PaddleOCR
import os
import random
from datetime import datetime
import numpy as np
import time

# 全局常量定义（目录路径）
STATIC_DIR = "log_image"  # 根目录

# 项目启动时检查并创建static根目录
if not os.path.exists(STATIC_DIR):
    os.makedirs(STATIC_DIR, exist_ok=True)
    print(f"已初始化创建static根目录：{os.path.abspath(STATIC_DIR)}")

# 创建FastAPI应用
app = FastAPI(root_path="/ocr_server")
router = APIRouter()


# 定义数据模型
class OCRResult(BaseModel):
    text: str
    position: List[List[float]]

class OCRResponse(BaseModel):
    result: dict
    msg: str
    code: int


# ===================== 全局初始化PaddleOCR实例 =====================
# 只在项目启动时初始化一次，
start_time = time.time()
ocr = PaddleOCR(use_gpu=True, gpu_id=0, use_angle_cls=True, lang="ch")
load_time = time.time() - start_time
print(f"PaddleOCR模型加载完成，耗时：{load_time:.2f}秒")


@router.post("/ocr", response_model=OCRResponse)
async def recognize_text(file: UploadFile = File(...)):
    try:
        # 生成唯一目录名：年月日时分秒-随机数（随机数取8位，降低重复概率）
        timestamp = datetime.now().strftime("%Y年%m月%d日-%H时%M分%S秒")
        random_num = random.randint(10000000, 99999999)  # 8位随机数
        unique_dir = f"{timestamp}-{random_num}"
        
        # 完整保存目录：static/唯一目录名
        save_dir = os.path.join(STATIC_DIR, unique_dir)
        
        # 创建当前请求的专属目录
        os.makedirs(save_dir, exist_ok=True)
        
        # 生成保存文件的路径（直接存放在专属目录下，无需再加时间戳）
        file_extension = os.path.splitext(file.filename)[1] if file.filename else ".jpg"
        # 文件名简化为"original" + 扩展名，因为目录已经唯一
        save_filename = f"original{file_extension}"
        file_path = os.path.join(save_dir, save_filename)
        
        # 保存上传的文件
        with open(file_path, "wb") as f:
            f.write(await file.read())

        # 使用OpenCV读取图片并获取长宽
        img = cv2.imread(file_path)
        if img is None:
            raise ValueError(f"无法读取图片文件，可能是格式错误或文件损坏：{file_path}")
        height, width = img.shape[:2]

        # ===================== 使用全局ocr实例 =====================
        # 执行OCR识别
        ocr_start_time = time.time()
        result = ocr.ocr(file_path, cls=True)
        ocr_cost_time = time.time() - ocr_start_time
        print(f"OCR识别完成，耗时：{ocr_cost_time:.4f}秒，输入的图纸名字：{file.filename}，保存到：{file_path}")


        # 提取文本和位置信息
        texts = []
        if result and result[0]:
            for line in result[0]:
                if line and len(line) > 1 and line[1]:
                    text = line[1][0]
                    position = line[0]
                    ocr_result = OCRResult(text=text, position=position)
                    texts.append(ocr_result)

        response_data = OCRResponse(
            result={
                "image_shape": [height, width], 
                "texts": texts,
                "storage_dir": unique_dir
            },
            msg="success",
            code=0
        )
        
        # === 新增：生成可视化图像 ===
        vis_img = img.copy()
        for ocr_res in texts:
            points = np.array(ocr_res.position, dtype=np.int32)
            cv2.polylines(vis_img, [points], isClosed=True, color=(0, 255, 0), thickness=2)
        
        vis_path = os.path.join(save_dir, "visualized.jpg")
        cv2.imwrite(vis_path, vis_img)

        # === 保存响应结果到 JSON 文件 ===
        import json
        result_json_path = os.path.join(save_dir, "result.json")
        with open(result_json_path, "w", encoding="utf-8") as f:
            json.dump(response_data.model_dump(), f, ensure_ascii=False, indent=2)

        return response_data


    except Exception as e:
        return OCRResponse(
            result={},
            msg=str(e),
            code=500
        )

# 将路由添加到应用
app.include_router(router)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7503)